from collections import Counter
from math import log

hmm_model = {i:Counter() for i in 'sbme'}

with open('dict.txt') as f:
    for line in f:
        lines = line.decode('utf-8').split(' ')
        if len(lines[0]) == 1:
            hmm_model['s'][lines[0]] += int(lines[1])
        else:
            hmm_model['b'][lines[0][0]] += int(lines[1])
            hmm_model['e'][lines[0][-1]] += int(lines[1])
            for m in lines[0][1:-1]:
                hmm_model['m'][m] += int(lines[1])

log_total = {i:log(sum(hmm_model[i].values())) for i in 'sbme'}

trans = {'ss':0.3,
         'sb':0.7,
         'bm':0.3,
         'be':0.7,
         'mm':0.3,
         'me':0.7,
         'es':0.3,
         'eb':0.7
         }

trans = {i:log(j) for i,j in trans.iteritems()}

def viterbi(nodes):
    paths = nodes[0]
    for l in range(1, len(nodes)):
        paths_ = paths
        paths = {}
        for i in nodes[l]:
            nows = {}
            for j in paths_:
                if j[-1]+i in trans:
                    nows[j+i]= paths_[j]+nodes[l][i]+trans[j[-1]+i]
            k = nows.values().index(max(nows.values()))
            paths[nows.keys()[k]] = nows.values()[k]
    return paths.keys()[paths.values().index(max(paths.values()))]

def hmm_cut(s):
    nodes = [{i:log(j[t]+1)-log_total[i] for i,j in hmm_model.iteritems()} for t in s]
    tags = viterbi(nodes)
    words = [s[0]]
    for i in range(1, len(s)):
        if tags[i] in ['b', 's']:
            words.append(s[i])
        else:
            words[-1] += s[i]
    return words
